import json
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# 读取Excel文件
df = pd.read_excel('dataclean_law.xlsx')


# 分句函数
def split_sentences(text):
    # 使用常见的句子结束标点符号进行分割
    sentence_delimiters = ['.', '?', '!', '；', '。', '？', '！']
    sentences = []
    current_sentence = ""
    for char in text:
        current_sentence += char
        if char in sentence_delimiters:
            sentences.append(current_sentence.strip())  # 将当前句子添加到句子列表中
            current_sentence = ""
    if current_sentence:
        sentences.append(current_sentence.strip())  # 将剩余的句子添加到句子列表中

    return sentences


# # 收集所有句子列表的大列表
# all_sentences = []
#
# # 遍历文件中的' fact '列并分句
# for text in df['fact']:
#     sentences = split_sentences(text)
#     all_sentences.extend(sentences)
#
# # 打印句子列表内容
# print(all_sentences)
#
# # 创建一个字典，键为 'sentences'，值为所有句子的列表
# data = {'sentences': all_sentences}
#
# # 导出为JSON文件（以UTF-8编码）
# with open('sentences.json', 'w', encoding='utf-8') as file:
#     json.dump(data, file, ensure_ascii=False)


# BM25算法函数
def bm25(query, sentences):
    tfidf_vectorizer = TfidfVectorizer()    # 初始化TF-IDF向量化器
    tfidf_matrix = tfidf_vectorizer.fit_transform([query] + sentences)  # 对query和分句文本进行TF-IDF向量化
    cosine_similarities = cosine_similarity(tfidf_matrix[0:1], tfidf_matrix[1:]).flatten()  # 计算余弦相似度
    output = {
        'query': query,  # 查询文本
        'sentences': sentences,  # 分句文本列表
        'cosine_similarities': cosine_similarities.tolist()  # 余弦相似度的计算结果转换为列表形式
    }

    output_json = json.dumps(output, indent=4, ensure_ascii=False)  # 将输出字典转换为JSON格式的文本（不转义非ASCII字符）

    with open('output.json', 'a', encoding='utf-8') as file:  # 以UTF-8编码方式写入文件（追加写入模式）
        file.write(output_json)  # 将JSON格式的文本写入文件

    return cosine_similarities  # 返回余弦相似度结果


# 筛选相似句子函数
def filter_similar_sentences(query, sentences):
    similarities = bm25(query, sentences)  # 使用BM25算法计算相似度，使用fact中的所有与query进行筛选
    # similarities = bm25(query, sentences[:3])  # 使用fact中的前三行与query进行筛选(也可修改为前n行，如下也如此)
    similar_sentences = []
    for i, similarity in enumerate(similarities):
        if similarity > 0.7 and i < 3:  # 如果相似度大于阈值且是开头的句子，则忽略该句子
            continue
        elif similarity > 0.7:  # 如果相似度大于阈值
            sentence = sentences[i-1]  # 获取相似的句子
            if sentence not in sentences[:3] and sentence not in similar_sentences:  # 如果句子不在开头的三句和已选的相似句子中
                similar_sentences.append(sentence)  # 将句子添加到相似句子列表中
    return similar_sentences


# 对每一行的fact列进行处理
# 提取开头的三句话并拼接成一个字符串
df['fact_fixed'] = df['fact'].apply(lambda x: ' '.join(split_sentences(x)[:3]))
# 对fact列中的内容进行分句
df['fact_sentences'] = df['fact'].apply(split_sentences)
# 筛选相似句子
df['similar_sentences'] = df.apply(lambda x: filter_similar_sentences(x['query'], x['fact_sentences']), axis=1)

# 组合开头的三句和筛选出的相似句子为新的“案件基本情况”内容
df['案件基本情况'] = df.apply(lambda x: ' '.join(split_sentences(x['fact_fixed']) + x['similar_sentences']), axis=1)

# 只保留id列和案件基本情况列
df = df[['id', '案件基本情况']]

# 导出文件
# 对fact列内的全部文本进行筛选
# df.to_csv('processed_data_all.csv', index=False)
# 对fact列内容选取前三句进行筛选
# df.to_csv('processed_data_3.csv', index=False)


# 检验完整性
# 读取两个CSV文件
df1 = pd.read_csv('processed_data_all.csv')
df2 = pd.read_csv('processed_data_3.csv')

# 设置索引为'id'列
df1.set_index('id', inplace=True)
df2.set_index('id', inplace=True)

# 获取两个DataFrame之间的共同'id'值
common_ids = df1.index.intersection(df2.index)

# 提取共同id对应的案件基本情况
col1 = df1.loc[common_ids, '案件基本情况'].astype(str)
col2 = df2.loc[common_ids, '案件基本情况'].astype(str)

# 初始化TF-IDF向量化器
vectorizer = TfidfVectorizer()
# 对第一个文件中的案件基本情况进行拟合并转换为TF-IDF向量
tfidf_matrix1 = vectorizer.fit_transform(col1)
# 对第二个文件中的案件基本情况进行转换为TF-IDF向量
tfidf_matrix2 = vectorizer.transform(col2)

# 计算每个id对应案件之间的相似度
similarity_values = cosine_similarity(tfidf_matrix1, tfidf_matrix2)

# 构建一个新的DataFrame，记录相似度和对应的id
similarity_df = pd.DataFrame(similarity_values, index=common_ids, columns=common_ids)

# 输出相似度
print("相似度矩阵:")
print(similarity_df)

# 导出相似度矩阵为CSV文件
similarity_df.to_csv('similarity_matrix.csv', index=True, header=True)

